#include "tcp_sender.hh"

#include "tcp_config.hh"
#include "tcp_segment.hh"
#include "wrapping_integers.hh"

#include <cstdint>
#include <cstdio>
#include <random>

// Dummy implementation of a TCP sender

// For Lab 3, please replace with a real implementation that passes the
// automated checks run by `make check_lab3`.

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;

//! \param[in] capacity the capacity of the outgoing byte stream
//! \param[in] retx_timeout the initial amount of time to wait before retransmitting the oldest outstanding segment
//! \param[in] fixed_isn the Initial Sequence Number to use, if set (otherwise uses a random ISN)
TCPSender::TCPSender(const size_t capacity, const uint16_t retx_timeout, const std::optional<WrappingInt32> fixed_isn)
    : _isn(fixed_isn.value_or(WrappingInt32{random_device()()}))
    , _initial_retransmission_timeout{retx_timeout}
    , _stream(capacity) {
    // set rto
    _timer.reset();
    _timer.rto() = _initial_retransmission_timeout;
}

uint64_t TCPSender::bytes_in_flight() const { return _next_seqno - _abso_ackno; }

void TCPSender::fill_window() {
    size_t win_size = max(_window_size, 1ul);

    while (!_fin && _next_seqno < _abso_ackno + win_size) {
        TCPSegment seg{};
        seg.header().seqno = wrap(_next_seqno, _isn);

        if (!_syn) {
            seg.header().syn = true;
            _syn = true;
        }

        size_t payload_len = min(win_size + _abso_ackno - _next_seqno - seg.header().syn, TCPConfig::MAX_PAYLOAD_SIZE);
        std::string data = _stream.read(payload_len);

        if (_stream.eof() && data.length() + _next_seqno + seg.header().syn + 1 <= _abso_ackno + win_size) {
            seg.header().fin = true;
            _fin = true;
        }
        seg.payload() = Buffer(std::move(data));

        if (seg.length_in_sequence_space() == 0) {
            break;
        }
        _next_seqno += seg.length_in_sequence_space();

        _segments_out.emplace(seg);
        _ost_segs.emplace(std::move(seg));

        if (!_timer.is_running()) {
            _timer.start();
        }
    }
}

//! \param ackno The remote receiver's ackno (acknowledgment number)
//! \param window_size The remote receiver's advertised window size
void TCPSender::ack_received(const WrappingInt32 ackno, const uint16_t window_size) {
    _window_size = window_size;

    uint64_t new_ackno = unwrap(ackno, _isn, _abso_ackno);
    if (new_ackno > _next_seqno) {
        return;
    }
    if (new_ackno > _abso_ackno) {
        _timer.rto() = _initial_retransmission_timeout;
        _timer.reset();
        _consec_rtx = 0;
    }

    _abso_ackno = max(_abso_ackno, new_ackno);
    while (!_ost_segs.empty()) {
        auto &seg = _ost_segs.front();
        uint64_t seqno = unwrap(seg.header().seqno, _isn, _abso_ackno);
        if (seqno + seg.length_in_sequence_space() <= _abso_ackno) {
            _ost_segs.pop();
        } else {
            break;
        }
    }
    if (_ost_segs.empty()) {
        _timer.stop();
    }
    fill_window();
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void TCPSender::tick(const size_t ms_since_last_tick) {
    _timer.update(ms_since_last_tick);
    if (_timer.is_expired()) {
        if (!_ost_segs.empty()) {
            _segments_out.emplace(_ost_segs.front());
            if (_window_size) {
                ++_consec_rtx;
                _timer.rto() <<= 1;
            }
        }
        _timer.start();
    }
}

unsigned int TCPSender::consecutive_retransmissions() const { return _consec_rtx; }

void TCPSender::send_empty_segment() {
    TCPSegment seg{};
    seg.header().seqno = next_seqno();
    _segments_out.emplace(std::move(seg));
}

void TCPSender::Timer::reset() { _total_ms = 0; }

bool TCPSender::Timer::is_expired() const {
    if (!_is_running) {
        return false;
    }
    if (_total_ms >= _rto) {
        return true;
    }
    return false;
}

void TCPSender::Timer::update(const uint64_t time_passed) {
    if (_is_running)
        _total_ms += time_passed;
}

bool TCPSender::Timer::is_running() const { return _is_running; }
void TCPSender::Timer::start() {
    _is_running = true;
    reset();
}
void TCPSender::Timer::stop() {
    _is_running = false;
    reset();
}
